package Data_Structure.Tree.Heap;

import java.util.*;

/**
 * @className: Test3
 * @author: 芃
 * @description: 堆相关OJ习题
 * @date: 2023/8/25 17:08
 */
public class Test3 {
    public static void main(String[] args) {
        Random random = new Random();
        int[] array = new int[1000];
        for (int i = 0; i < array.length; i++) {
            array[i] = random.nextInt(1000);
        }
        //如果是要返回第K大或者第K小的元素，返回的数组首元素即为结果
        int[] ret = topK(array,5);
        System.out.println(Arrays.toString(ret));
        int[] ret2 = topK2(array,5);
        System.out.println(Arrays.toString(ret2));
    }
    /**
     * @param array:
     * @return int
     * @author: zyp19
     * @description: 求数组当中前K个最大的元素
     * @date: 2023/8/25 17:09
     */
    public static int[] topK(int[] array,int k){
        //①、创建一个大小为K的小根堆
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o1 - o2;
            }
        });
        //②、前K个元素放到小根堆中
        for (int i = 0; i < k; i++) {
            priorityQueue.offer(array[i]);
        }
        //③、将剩下的元素和堆顶元素进行比较，如果大于堆顶元素，就将堆顶元素弹出，放入新的元素
        for (int i = k; i < array.length; i++) {
            if (priorityQueue.peek() < array[i]){
                priorityQueue.poll();
                priorityQueue.offer(array[i]);
            }
        }
        //④、将堆里面的元素弹出放入到新的容器中，因为是小根堆，所以默认是从小到大弹出
        int[] ret = new int[k];
        int index = 0;
        while (!priorityQueue.isEmpty()){
            ret[index] = priorityQueue.poll();
            index++;
        }
        //⑤、根据情况而定，如果是需要从大到小返回，逆置即可
        return ret;
    }
    /**
     * @param array:
     * @param k:
     * @return int
     * @author: zyp19
     * @description: 求数组当中前K个最小的元素
     * @date: 2023/8/25 17:19
     */
    public static int[] topK2(int[] array,int k){
        //①、创建一个大小为K的大根堆
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });
        //②、前K个元素放到小根堆中
        for (int i = 0; i < k; i++) {
            priorityQueue.offer(array[i]);
        }
        //③、将剩下的元素和堆顶元素进行比较，如果小于堆顶元素，就将堆顶元素弹出，放入新的元素
        for (int i = k; i < array.length; i++) {
            if (priorityQueue.peek() > array[i]){
                priorityQueue.poll();
                priorityQueue.offer(array[i]);
            }
        }
        //④、将堆里面的元素弹出放入到新的容器中，因为是大根堆，所以默认是从大到小弹出
        int[] ret = new int[k];
        int index = 0;
        while (!priorityQueue.isEmpty()){
            ret[index] = priorityQueue.poll();
            index++;
        }
        //⑤、根据情况而定，如果是需要从小到大返回，逆置即可
        return ret;
    }
    /**
     * @param nums1:
     * @param nums2:
     * @param k:
     * @return List<List<Integer>>
     * @author: zyp19
     * @description: 查找和最小的k对数字
     * 什么意思？给定两个数组，在两个数组中任取一个数字，组成一个数对，求和最小的前K对数字
     * @date: 2023/8/25 17:22
     */
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> ret = new ArrayList<>();
        PriorityQueue<List<Integer>> maxHeap = new PriorityQueue<>(k, new Comparator<List<Integer>>() {
            @Override
            public int compare(List<Integer> o1, List<Integer> o2) {
                int val1 = o1.get(0) + o1.get(1);
                int val2 = o2.get(0) + o2.get(1);
                return val2 - val1;
            }
        });
        int index = 0;
        for (int value1 : nums1) {
            for (int value2 : nums2) {
                List<Integer> list = new ArrayList<>(2);
                list.add(value1);
                list.add(value2);
                if (index < k){
                    maxHeap.offer(list);
                }else {
                    int v1 = list.get(0) + list.get(1);
                    int v2 = maxHeap.peek().get(0) + maxHeap.peek().get(1);
                    if (v1 < v2){
                        maxHeap.poll();
                        maxHeap.offer(list);
                    }else{
                        break;
                    }
                }
                index++;
            }
        }
        while (!maxHeap.isEmpty()){
            List<Integer> res = maxHeap.poll();
            ret.add(res);
        }
        return ret;
    }
}
